/**
 * @file main.cpp
 *
 * Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <cstdint>
#include <iostream>

#include "acl/acl.h"
#include "common.h"
#include "op_runner.h"

bool g_isDevice = false;
int deviceId = 0;

OperatorDesc CreateOpDesc()
{
    // 根据实际情况请修改shape、datatype、format
    std::vector<int64_t> shape{1093};  // 一维形状
    aclDataType dataType = ACL_FLOAT;   // float32
    aclFormat format = ACL_FORMAT_ND;   // 多维格式（兼容一维）

    OperatorDesc opDesc;
    // 有几个输入就AddInputTensorDesc几次
    opDesc.AddInputTensorDesc(dataType, shape.size(), shape.data(), format);
    opDesc.AddInputTensorDesc(dataType, shape.size(), shape.data(), format);
    opDesc.AddInputTensorDesc(dataType, shape.size(), shape.data(), format);
    opDesc.AddInputTensorDesc(dataType, shape.size(), shape.data(), format);

    opDesc.AddOutputTensorDesc(dataType, shape.size(), shape.data(), format);
    opDesc.AddOutputTensorDesc(dataType, shape.size(), shape.data(), format);
    // opDesc.AddAttr("momentum", 0.9f);  // 对应索引 8
    // opDesc.AddAttr("epsilon", 1e-5f);   // 对应索引 9
    return opDesc;
}

bool SetInputData(OpRunner &runner)
{
    size_t fileSize = 0;
    //确定输入个数
    ReadFile("../input/local_mean.bin", fileSize, runner.GetInputBuffer<void>(0), runner.GetInputSize(0));
    ReadFile("../input/local_var.bin", fileSize, runner.GetInputBuffer<void>(1), runner.GetInputSize(1));
    ReadFile("../input/global_mean.bin", fileSize, runner.GetInputBuffer<void>(2), runner.GetInputSize(2));
    ReadFile("../input/global_var.bin", fileSize, runner.GetInputBuffer<void>(3), runner.GetInputSize(3));

    // 打印输入数据的前几个值，确认输入正确
    auto* local_mean = runner.GetInputBuffer<float>(0);
    INFO_LOG("First 4 values of local_mean: %f, %f, %f, %f",
             local_mean[0], local_mean[1], local_mean[2], local_mean[3]);
    auto* local_var = runner.GetInputBuffer<float>(1);
    INFO_LOG("First 4 values of local_var: %f, %f, %f, %f",
             local_var[0], local_var[1], local_var[2], local_var[3]);
    auto* global_mean = runner.GetInputBuffer<float>(2);
    INFO_LOG("First 4 values of global_mean: %f, %f, %f, %f",
             global_mean[0], global_mean[1], global_mean[2], global_mean[3]);
    auto* global_var = runner.GetInputBuffer<float>(3);
    INFO_LOG("First 4 values of global_var: %f, %f, %f, %f",
             global_var[0], global_var[1], global_var[2], global_var[3]);                          
    INFO_LOG("Set input success");
    return true;
}



bool ProcessOutputData(OpRunner &runner)
{   
    //确定输出个数
    WriteFile("../output/updated_mean.bin", runner.GetOutputBuffer<void>(0), runner.GetOutputSize(0));
    WriteFile("../output/updated_var.bin", runner.GetOutputBuffer<void>(1), runner.GetOutputSize(1));
    INFO_LOG("Write output success");
    return true;
}

void DestroyResource()
{
    bool flag = false;
    if (aclrtResetDevice(deviceId) != ACL_SUCCESS) {
        ERROR_LOG("Reset device %d failed", deviceId);
        flag = true;
    }
    INFO_LOG("Reset Device success");
    if (aclFinalize() != ACL_SUCCESS) {
        ERROR_LOG("Finalize acl failed");
        flag = true;
    }
    if (flag) {
        ERROR_LOG("Destroy resource failed");
    } else {
        INFO_LOG("Destroy resource success");
    }
}

bool InitResource()
{
    std::string output = "../output";
    if (access(output.c_str(), 0) == -1) {
        int ret = mkdir(output.c_str(), 0700);
        if (ret == 0) {
            INFO_LOG("Make output directory successfully");
        } else {
            ERROR_LOG("Make output directory fail");
            return false;
        }
    }

    // acl.json is dump or profiling config file
    if (aclInit("../scripts/acl.json") != ACL_SUCCESS) {
        ERROR_LOG("acl init failed");
        return false;
    }

    if (aclrtSetDevice(deviceId) != ACL_SUCCESS) {
        ERROR_LOG("Set device failed. deviceId is %d", deviceId);
        (void)aclFinalize();
        return false;
    }
    INFO_LOG("Set device[%d] success", deviceId);

    // runMode is ACL_HOST which represents app is running in host
    // runMode is ACL_DEVICE which represents app is running in device
    aclrtRunMode runMode;
    if (aclrtGetRunMode(&runMode) != ACL_SUCCESS) {
        ERROR_LOG("Get run mode failed");
        DestroyResource();
        return false;
    }
    g_isDevice = (runMode == ACL_DEVICE);
    INFO_LOG("Get RunMode[%d] success", runMode);

    return true;
}

bool RunOp()
{
    // create op desc
    OperatorDesc opDesc = CreateOpDesc();

    // create Runner
    OpRunner opRunner(&opDesc);
    if (!opRunner.Init()) {
        ERROR_LOG("Init OpRunner failed");
        return false;
    }
    INFO_LOG("OpRunner initialized successfully");

    // Load inputs
    if (!SetInputData(opRunner)) {
        ERROR_LOG("Set input data failed");
        return false;
    }
    // Run op
    if (!opRunner.RunOp()) {
        ERROR_LOG("Run op failed");
        return false;
    }
    INFO_LOG("Run op success");

    // ===== 新增调试打印 =====
    for (int i = 0; i < 2; ++i) {  // 输出 buffer 数量，一般 Mean/Var 两个
        const float* buf = opRunner.GetOutputBuffer<float>(i);
        size_t buf_size = opRunner.GetOutputSize(i);

        INFO_LOG("Output buffer[%d] address: %p", i, buf);
        INFO_LOG("Output buffer[%d] size: %zu", i, buf_size);

        if (buf_size >= 4 * sizeof(float)) {
            INFO_LOG("First 4 floats of output buffer[%d]: %f, %f, %f, %f",
                     i, buf[0], buf[1], buf[2], buf[3]);
        } else {
            INFO_LOG("Output buffer[%d] too small to display first 4 floats", i);
        }
    }
    // ===== 调试打印结束 =====

    // process output data
    if (!ProcessOutputData(opRunner)) {
        ERROR_LOG("Process output data failed");
        return false;
    }

    INFO_LOG("Run op success");
    return true;
}

int main(int argc, char **argv)
{
    if (!InitResource()) {
        ERROR_LOG("Init resource failed");
        return FAILED;
    }
    INFO_LOG("Init resource success");

    if (!RunOp()) {
        DestroyResource();
        return FAILED;
    }

    DestroyResource();

    return SUCCESS;
}
